from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa


# 1. 生成 RSA 密钥对
def generate_keys():
    private_key = rsa.generate_private_key(
        public_exponent=65537, key_size=2048, backend=default_backend()
    )
    public_key = private_key.public_key()
    return private_key, public_key


# 2. 使用私钥进行签名
def sign_message(private_key, message):
    signature = private_key.sign(
        message.encode("utf-8"),  # 将消息编码为字节
        padding.PSS(
            mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH
        ),
        hashes.SHA256(),
    )
    return signature


# 3. 使用公钥验证签名
def verify_signature(public_key, message, signature):
    try:
        public_key.verify(
            signature,
            message.encode("utf-8"),
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256(),
        )
        print("签名验证成功！")
        return True
    except Exception as e:
        print("签名验证失败:", e)
        return False


# 主程序
if __name__ == "__main__":
    # 生成密钥对
    private_key, public_key = generate_keys()

    # 要签名的消息
    message = "这是一个需要签名的消息。"

    # 进行签名
    signature = sign_message(private_key, message)
    print("签名:", signature.hex())

    # 验证签名
    is_valid = verify_signature(public_key, message, signature)

    # 尝试验证一个错误的签名
    wrong_message = "这是一个错误的消息。"
    is_valid_wrong = verify_signature(public_key, wrong_message, signature)
